#!/usr/bin/env python

import numpy as np
import pylab as pl	
from pylab import NaN
import os.path as osp
import sys
import matplotlib
from matplotlib.font_manager import FontProperties as FP

def ProbabilityMatrix(gen_size,q):
	# Generate a probability matrix for EEP generation
	P=np.zeros(shape=(gen_size+1,gen_size+1))
	for i in range(gen_size+1):
		entry_val=1./(q**(gen_size-i))
		P[i,i]=entry_val
		if i<gen_size:
			P[i+1,i]=1-entry_val
	return P
	

def DecodingProbability(P,pkts_range):
	# Calculate decoding probability 

	# initial state vector
	s=np.zeros(shape=(gen_size+1))
	s[0]=1
	
	# Save decoding probabilities here
	decod_prob=[]

	# For each packet sent calc. prob
	for nb_pkts in pkts_range:
			decod_prob_tmp=np.dot(np.linalg.matrix_power(P,nb_pkts),s)
			decod_prob.append(decod_prob_tmp[-1])

	return decod_prob

if __name__ == '__main__':

	# Fixes font size
	fp = FP()
	fp.set_size('small')

	# Use latex formatting in strings
	matplotlib.rc('text', usetex=True) 

	# Prepare figure
	fig = pl.figure(figsize=[10,4]) # Setting aspect ratio
	ax1 = fig.add_subplot(111) # left plot

	# Parameters
	gen_size=50 # generation size
	q=[2**1,2**8,2**16] # fields to plot dec probs			

	# Generate plotting data sets
	P1=ProbabilityMatrix(gen_size,q[0]) # probability matrix 1
	P2=ProbabilityMatrix(gen_size,q[1]) # probability matrix 2
	P3=ProbabilityMatrix(gen_size,q[2]) # probability matrix 3
#	P4=ProbabilityMatrix(gen_size,2**4) # probability matrix 4

	# Which packets should we calculate probability for
	pkts_range=range(0,3*gen_size)

	# Calculate decoding probabilities
	dec_prob1=DecodingProbability(P1,pkts_range)
	dec_prob2=DecodingProbability(P2,pkts_range)
	dec_prob3=DecodingProbability(P3,pkts_range)
#	dec_prob4=DecodingProbability(P4,pkts_range)

	# Layout fix - we do not want 0's before 'g' transmitted packets
	# Crappy solution, why is alternative not working?
##	dec_prob1[0:(gen_size-1)]=np.NaN
##	dec_prob2[0:gen_size-1]=np.NaN
##	dec_prob3[0:gen_size-1]=np.NaN
	for i in range(len(dec_prob1)):
		if i<gen_size:
			dec_prob1[i]=np.NaN
			dec_prob2[i]=np.NaN
			dec_prob3[i]=np.NaN

	# Fill plt
	ax1.plot(pkts_range,dec_prob1,color='black',label='FF(2), g='+str(gen_size),marker='s',mfc='none')
	ax1.plot(pkts_range,dec_prob2,label='FF($2^8$), g='+str(gen_size),marker='o',color='black',mfc='none')
	ax1.plot(pkts_range,dec_prob3,label='FF($2^{16}$), g='+str(gen_size),marker='^',color='black',mfc='none')
#	ax1.plot(pkts_range,dec_prob4,label='GF($2^{4}$), g='+str(gen_size)',marker='*')

	# Plot annotation
	pl.xlabel('Number of received linear combinations [-]')
	pl.ylabel('Probability of decoding the source data [-]')
	pl.legend(loc='lower right',prop=fp)
	pl.grid('on')
	pl.xticks(np.arange(0,3*gen_size,1))
	pl.yticks(np.arange(0,1.2,0.1))
	pl.xlim(gen_size-1,gen_size+10)
	pl.ylim(0,1.1) # probability

#	# Save figure
	svndir = osp.abspath('../../rapport/figs')
	pl.savefig(svndir+'/'+'decoding_prob'+'.eps')




